#include <iostream>
#include <vector>
#include <atomic>
#include <thread>
#include <mutex>
#include <bitset>
#include <boost/container/static_vector.hpp>

std::mutex g_mutex;

static constexpr size_t QUEUE_SIZE = 100000;
static constexpr size_t TOTAL_COUNT = 1000000;
static constexpr size_t push_thread_cnt = 2;
static constexpr size_t pop_thread_cnt = 5;

template <typename T>
class lock_free_ring
{
public:
    lock_free_ring(size_t size)
        : cons_tail_(size - 1), cons_head_(size - 1), data_(size), size_(size)
    {
    }

    bool push(const T &value)
    {
        auto prod_head = prod_head_.load(std::memory_order_relaxed);
        size_t new_prod_head = 0;
        size_t cons_tail = 0;

        do
        {
            std::atomic_thread_fence(std::memory_order_acquire);
            cons_tail = cons_tail_.load(std::memory_order_acquire);
            new_prod_head = (prod_head + 1) % size_;
            if (new_prod_head == cons_tail || prod_head == cons_tail)
                return false;
        } while (!prod_head_.compare_exchange_strong(prod_head, new_prod_head, std::memory_order_relaxed, std::memory_order_relaxed));

        data_[prod_head] = value;

        while (prod_tail_.load(std::memory_order_relaxed) != prod_head)
            __builtin_ia32_pause();

        prod_tail_.store(new_prod_head, std::memory_order_release);
        return true;
    }

    bool pop(T &value)
    {
        size_t new_cons_head = 0;
        auto cons_head = cons_head_.load(std::memory_order_relaxed);
        size_t prod_tail = 0;

        do
        {
            std::atomic_thread_fence(std::memory_order_acquire);
            prod_tail = prod_tail_.load(std::memory_order_acquire);
            new_cons_head = (cons_head + 1) % size_;
            if (new_cons_head == prod_tail)
                return false;
        } while (!cons_head_.compare_exchange_strong(cons_head, new_cons_head, std::memory_order_relaxed, std::memory_order_relaxed));

        value = data_[new_cons_head];

        while (cons_tail_.load(std::memory_order_relaxed) != cons_head)
            __builtin_ia32_pause();

        cons_tail_.store(new_cons_head, std::memory_order_release);
        return true;
    }

    void print_info()
    {
        std::cout << "prod_tail_:" << prod_tail_.load(std::memory_order_relaxed)
                  << ", prod_head_:" << prod_head_.load(std::memory_order_relaxed)
                  << ", cons_tail_:" << cons_tail_.load(std::memory_order_relaxed)
                  << ", cons_head_:" << cons_head_.load(std::memory_order_relaxed)
                  << ", size_:" << size_ << std::endl;

        std::cout << "data: [";
        for (size_t i = 0; i < size_; ++i)
        {
            std::cout << data_[i] << ",";
            if (i % 20 == 0)
            {
                std::cout << std::endl;
            }
        }
        std::cout << "]" << std::endl;
    }

private:
    std::atomic<size_t> prod_tail_{0};
    std::atomic<size_t> prod_head_{0};
    char cache_line0_[0] __attribute__((aligned(64)));
    std::atomic<size_t> cons_tail_{0};
    std::atomic<size_t> cons_head_{0};
    char cache_line1_[0] __attribute__((aligned(64)));
    std::vector<T> data_{};
    size_t size_{0};
};

void push_thread(lock_free_ring<size_t> &lrq, size_t id)
{
    size_t i = id * (TOTAL_COUNT / push_thread_cnt);
    size_t target_cnt = i + (TOTAL_COUNT / push_thread_cnt);
    for (; i < target_cnt; ++i)
    {
        while (!lrq.push(i))
        {
            __builtin_ia32_pause();
        }
    }
}

void pop_thread(lock_free_ring<size_t> &lrq, std::bitset<TOTAL_COUNT> &data, volatile bool &stop)
{
    while (true)
    {
        size_t i = 0;
        if (lrq.pop(i))
        {
            data.set(i);
        }
        else
        {
            __builtin_ia32_pause();
            if (stop)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
                for (size_t i = 0; i < 1000; ++i)
                {
                    if (lrq.pop(i))
                    {
                        data.set(i);
                    }
                }
                break;
            }
        }
    }
}

void print_set(const std::string &title, const std::bitset<TOTAL_COUNT> &data)
{
    std::cout << title << ": [";

    for (size_t i = 0; i < TOTAL_COUNT; ++i)
    {
        if (data.test(i))
        {
            std::cout << i << ",";
            if (i % 20 == 0)
                std::cout << std::endl;
        }
    }

    std::cout << "]" << std::endl;
}

int main()
{
    lock_free_ring<size_t> lrq(QUEUE_SIZE);

    std::bitset<TOTAL_COUNT> expect;
    for (size_t i = 0; i < TOTAL_COUNT; ++i)
        expect.set(i, true);

    size_t succ_cnt = 0;

    while (true)
    {
        boost::container::static_vector<std::thread, push_thread_cnt> push_threads;
        volatile bool stop = false;
        for (size_t i = 0; i < push_thread_cnt; ++i)
            push_threads.push_back(std::thread(push_thread, std::ref(lrq), i));

        boost::container::static_vector<std::thread, pop_thread_cnt> pop_threads;
        std::bitset<TOTAL_COUNT> data[pop_thread_cnt];
        for (size_t i = 0; i < pop_thread_cnt; ++i)
            pop_threads.push_back(std::thread(pop_thread, std::ref(lrq), std::ref(data[i]), std::ref(stop)));

        for (size_t i = 0; i < push_thread_cnt; ++i)
            push_threads[i].join();

        stop = true;
        for (size_t i = 0; i < pop_thread_cnt; ++i)
            pop_threads[i].join();

        std::bitset<TOTAL_COUNT> result;
        for (auto &d : data)
        {
            // 判断是否有重复的数据
            auto tmp = d & result;
            if (tmp.any())
            {
                std::cout << "data duplicated" << std::endl;

                // get object index
                auto const index = &d - &data[0];
                std::cout << "duplicated at: " << index << std::endl;

                // print duplicated data
                print_set("duplicate", tmp);
                return -1;
            }

            result |= d;
        }

        if (expect == result)
        {
            succ_cnt++;
            std::cout << "pass " << succ_cnt << '\r' << std::flush;
        }
        else
        {
            std::cout << std::endl;
            std::cout << "expect_count: " << expect.size() << ", result_count: " << result.size() << std::endl;
            std::cout << "fail" << std::endl;
            lrq.print_info();
            // print_set("expect", expect);
            // print_set("result", result);
            std::bitset<TOTAL_COUNT> diff;

            // diff bitset
            for (size_t i = 0; i < TOTAL_COUNT; ++i)
                if (expect.test(i) != result.test(i))
                    diff.set(i, true);

            print_set("diff", diff);
            return -1;
        }
    }
    return 0;
}